from asyncore import write
import random
import queue
import numpy as np
import math
import cv2

# seed
seed = 0
r_max = 1e9+7

# area property
length = int(2e3)    # current area size
smoothness = 25 # current area corner smoothness

# dem graph
MAX = int(2e3+3)
dem = np.zeros([MAX, MAX])
vis = np.zeros([MAX, MAX])
moon = np.zeros([MAX, MAX])

#
min_crater_d = 10
min_rock_d = 1.5
resolution = 0.5
h1_d= [[0.23,0.17], [0.11,0.25], [0.19,0.13]] # h1_d[i][0] for lower_limit, h1_d[i][1] for upper_limit
h2_d= [[0.022,0.016], [0.008,0.06], [0.045,0.03]] # ditto

class area:
    def __init__(self, xs, xt, ys, yt, s):
        self.xs = int(xs)
        self.xt = int(xt)
        self.ys = int(ys)
        self.yt = int(yt)
        self.s = s    # smoothness
        
def randomf():
    global r_max
    return random.uniform(0, r_max)

def randomi():
    global r_max
    return random.randint(1, r_max)

# the crater height of current position(x, y)
def crater(d, h1, h2, x, y):
    x = math.sqrt(sq(x) + sq(y))
    rd = randomf() % 2 + 1
    if (x + rd <= (d/2)):
        res = 4.0 * h1 / sq(d) * sq(x) + h2 - h1
    else:
        res = -4.0 * h2 / sq(d) * sq(x) + 4.0 * h2 / d * x
    if (x + rd > (d/2) and res < 0):
        return 0.0
    return res

# the rock height of current position(x, y)
def rock(d, h, x, y):
    res = h - 4.0 * h * (sq(x) + sq(y)) / sq(d)
    if (res < 0):
        return 0.0
    return res

def sq(x):
    return x*x

def seed_set():
    global seed
    if seed is not None:
        random.seed(seed)
    else:
        seed = random.randint(0, 1000)
        random.seed(seed)

# function of moon crater diameter to number
def cd2n(d):
    global length
    a = sq(length)   # moon area
    if d <= 40.0:
        return int(pow(10, (-2*math.log10(d)-1)) * (a / 4.0))
    else:
        return int(pow(10, (-3*math.log10(d)+0.62)) * (a / 4.0))
    
# function of number of moon crater to diameter
def cn2d(n):
    global length
    nn = n / sq(length/2)
    # print("{} {}".format(n, nn))
    if (n < 65):
        return pow(10, (math.log10(nn)-0.62)/-3.0)
    else:
        return pow(10, (math.log10(nn)+1.0)/-2.0)

def rd2n(d):
    global length
    a = sq(length)
    return int(pow(10, -2.589*math.log10(d)-3.3955) * a/4.0)

def rn2d(n):
    global length
    nn = n / sq(length/2)
    return pow(10, (math.log10(nn)+3.3955)/-2.589)

def fill_crater(d, h1, h2, xs, xt, ys, yt, xm, ym):
    global length, resolution, dem, vis
    '''
    d:  diameter of crater
    h1: inner height of crater
    h2: edge height of crater
    xs: x-r
    xt: x+r
    ys: y-r
    yt: y+r
    xm: x
    ym: y
    '''
    if (xs < 0): xs = 0
    if (xt >= length): xt = length - 1
    if (ys < 0): ys = 0
    if (yt >= length): yt = length - 1
    for i in range(xs, xt+1):
        for j in range(ys, yt+1):
            xx = (i - xm) * resolution
            yy = (j - ym) * resolution
            dem[i][j] += crater(d,h1,h2,xx,yy)
            vis[i][j] = 1

def fill_rock(d, h, xs, xt, ys, yt, xm, ym):
    global length, resolution, dem
    if (xs < 0): xs = 0
    if (xt >= length): xt = length - 1
    if (ys < 0): ys = 0
    if (yt >= length): yt = length - 1
    for i in range(xs, xt+1):
        for j in range(ys, yt+1):
            xx = (i - xm) * resolution
            yy = (j - ym) * resolution
            dem[i][j] += rock(d, h, xx, yy)

def terrain_generation():
    global length, smoothness, dem
    q = queue.Queue()
    q.put(area(0, length-1.0, 0, length-1.0, smoothness))
    dem[0][0] = dem[0][length-1] = dem[length-1][0] = dem[length][length] = 10.0
    while not q.empty():
        cur = q.get()    # include pop
        if (cur.xt - cur.xs < 2):
            continue
        xs = cur.xs
        xt = cur.xt
        ys = cur.ys
        yt = cur.yt
        mx = int((xs + xt) / 2)
        my = int((ys + yt) / 2)
        avg = (dem[xs][ys] + dem[xs][yt] + dem[xt][ys] + dem[xt][yt]) / 4.0 
        rd = randomf() % (max(cur.s * 2000.0, 1.0)) / 1000.0 - cur.s
        dem[mx][my] = avg + rd
        rd = randomf() % (max(cur.s * 2000.0, 1.0)) / 1000.0 - cur.s
        dem[mx][ys] = avg + rd
        rd = randomf() % (max(cur.s * 2000.0, 1.0)) / 1000.0 - cur.s
        dem[mx][yt] = avg + rd
        rd = randomf() % (max(cur.s * 2000.0, 1.0)) / 1000.0 - cur.s
        dem[xs][my] = avg + rd
        rd = randomf() % (max(cur.s * 2000.0, 1.0)) / 1000.0 - cur.s
        dem[xt][my] = avg + rd
        nxts = cur.s / 2.0
        q.put(area(xs, mx, ys, my, nxts))
        q.put(area(mx, xt, ys, my, nxts))
        q.put(area(xs, mx, my, yt, nxts))
        q.put(area(mx, xt, my, yt, nxts))

def crater_generation():
    global length, resolution, min_crater_d, h1_d, h2_d
    noc = cd2n(min_crater_d)  # number of crater
    print("number of crater is : {}".format(noc))
    minn = 1e9+7
    for i in range(noc):
        num = randomi() % noc + 1
        d = cn2d(num)
        minn = min(minn, d)
        t = randomi() % 3
        h1 = h1_d[t][0] + (h1_d[t][1] - h1_d[t][0]) * (randomf()%10+1) / 10.0
        h1 *= d
        h2 = h2_d[t][0] + (h2_d[t][1] - h2_d[t][0]) * (randomf()%10+1) / 10.0
        h2 *= d
        r = int(d / resolution)
        x = randomi() % length
        y = randomi() % length
        while(vis[x][y]):
            x = randomi() % length
            y = randomi() % length
        fill_crater(d, h1, h2, x-r, x+r, y-r, y+r, x, y)

def rock_generation():
    global length, resolution, min_rock_d
    nor = rd2n(min_rock_d)
    print("number of rock is : {}".format(nor))
    for i in range(nor):
        num = randomi() % nor + 1
        d = rn2d(num)
        h = (randomi() % 50 + 50) / 100.0 * d
        r = int((d / resolution) / 2)
        x = randomi() % length
        y = randomi() % length
        fill_rock(d, h, x-r, x+r, y-r, y+r, x, y)

def map_generation():
    global dem
    l = int(length / 4)
    with open("map.txt", 'w') as f:
        for i in range(l):
            for j in range(l):
                si = i * 4
                ti = j * 4
                sj = i * 4 + 4
                tj = j * 4 + 4
                for m in range(si, ti):
                    for n in range(sj, tj):
                        f.write("%.5f %.5f %.5f" % (m*0.5, n*0.5, dem[m][n]))
    cv2.imshow(dem)
    cv2.waitKey()
                  
    
def main():
    seed_set() # set random seed
    
    terrain_generation() # diamond-square algorithm
    
    crater_generation()
    
    rock_generation()
    
    map_generation()
    
if __name__ == "__main__":
    main()